import os
from Bio import SeqIO
from collections import defaultdict, Counter

def count_gene_annotations(genbank_file, gene_annotations):
    with open(genbank_file, "r") as file:
        for record in SeqIO.parse(file, "genbank"):
            for feature in record.features:
                if feature.type == "gene":
                    gene_name = feature.qualifiers.get("gene", ["unknown"])[0]
                    gene_annotations[gene_name]["count"] += 1
                    gene_annotations[gene_name]["samples"].add(genbank_file)
    return gene_annotations

def write_annotations_to_file(gene_annotations, output_file):
    with open(output_file, "w") as file:
        for gene, data in gene_annotations.items():
            sample_count = len(data["samples"])
            file.write(f"Gene: {gene}, Count: {data['count']}, Sample Files: {sample_count}\n")

def process_all_genbank_files():
    gene_annotations = defaultdict(lambda: {"count": 0, "samples": set()})
    
    for filename in os.listdir("."):
        if filename.endswith(".gb"):
            gene_annotations = count_gene_annotations(filename, gene_annotations)
    
    output_file = "combined_gene_annotations.txt"
    write_annotations_to_file(gene_annotations, output_file)
    print(f"Processed all .gb files and wrote results to {output_file}")

process_all_genbank_files()